# -*- coding: utf-8 -*-
"""
Created on Thu Nov 23 15:06:44 2023

solve free boundary equilibrium 

@author: xiangguchn
"""

import numpy as np
import math
import matplotlib.pyplot as plt

from scipy import special

import multiprocessing as mp
from calculate_green_function import green_function as GF

class free_boundary_gs():
    
    
    def __init__(self):
        # self.machine_name = settings.MACHINE_NAME
        # self.data_path = settings.DATA_PATH
        # self.machine_file = settings.MACHINE_FILE
        self.machine_file = 'C:\\Users\\xiang\\Documents\\XINAO\\AI\\SEPY\\data\\EXL50U\\machine.xlsx'
        # self.machine_file = 'C:\\sepy\\data\\EXL50U\\machine.xlsx'
        self.cmu = np.pi*2.0e-7;
        # self.num_cores = int(mp.cpu_count())
        eq = np.load('C:\\Users\\xiang\\Documents\\XINAO\\AI\\SEPY\\data\\EXL50U\\green_function.npy', 
                      allow_pickle=True).item()
        
        # eq = np.load('C:\\sepy\\data\\EXL50U\\green_function.npy', 
        #                 allow_pickle=True).item()
        
        # fgs dictionary is used for free gs solve
        self.fgs = {}
        
        # default region for finding x point 
        self.fgs['index_xpoint_up_x']   = 16
        self.fgs['index_xpoint_up_y']   = 20
        self.fgs['index_xpoint_down_x'] = 113
        self.fgs['index_xpoint_down_y'] = 20
        self.fgs['index_xpoint_height'] = 20
        self.fgs['index_xpoint_width']  = 20
        
        
        self.eq = eq
                
        self.num_cores = int(mp.cpu_count())
        
        
    def get_green_function(self):
        
        green = {}
        green['mi_plasma_plasma'] = np.zeros(129*129)
        green['mi_plasma_pf'] = np.zeros(shape=(129*129, 15))
        
        self.green = green
        return green
        
    
    
    def identity_plasma_init(self, flux_plasma, flux_limiter, init_factor=10, gapx=0):
        """
        find LCFS of plasma region with flux_plasma and flux_limiter. 
        judge plasma is limiter or divertor configuration with inductance on limiter

        input:
            option name         data_type    size     unit    meaning
            1      flux_plasma  array        (nr*nz,1)        flux of plasma induced by plasma and pfs
            1      flux_limiter array        (nlimiter,1)     flux of limiter induced by plasma and pfs
            0      isplot       int          1                is plot or not. 

        output:
            option  name            data_type   size    unit    meaning
            1      flux_plasma      array       (nr*nz,1)       flux of plasma induced by plasma and pfs
            1      flux_plasma_norm array       (nr*nz,1)       flux of plasma normalized
            1      index            array       (nr*nz,1)       index of plasma: 1 means inside LCFS; 0 means outside LCFS
            
        reference:
            [1]

        Developed by xiangguchn 2023/12/25
        """
                
        # plasma and limiter region
        x_plasma = self.eq['green']['grid']['plasma']['x']
        y_plasma = self.eq['green']['grid']['plasma']['y']
        
        x_limiter = self.eq['machine']['limiter']['x_limiter']
        y_limiter = self.eq['machine']['limiter']['y_limiter']
        # num_limiter = len(x_limiter)
        # num_init = len(x_init)
        
        if gapx==0:
            gapx = (x_plasma[1]-x_plasma[0])/2
        
        
        v_step = 1e-3 * init_factor
        ok_step = 5e-6 * init_factor
        
        my_delay = 0.5
        
        max_flux_plasma = max(flux_plasma)
        index_max_flux_plasma = np.argmax(flux_plasma)
        
        
        # plasma and limiter region
        
        # position higher than 1m should be cancelled
        y_top_limite = 1
        x_limiter_2 = x_limiter[abs(y_limiter)<y_top_limite]
        y_limiter_2 = y_limiter[abs(y_limiter)<y_top_limite]
        flux_limiter_2 = flux_limiter[abs(y_limiter)<y_top_limite]
        
        
        # find the max flux at limiter
        max_flux_limiter = max(flux_limiter_2)
        index_flux_limiter = np.argmax(flux_limiter_2)
        
        
        # find LCFS
        d_flux1 = max_flux_limiter
        d_flux0 = max_flux_plasma
        d_flux = d_flux1
        # plt.figure()
        # self.plot_flux(flux_plasma)
        while abs(d_flux0-d_flux1) > ok_step/10:
            d_flux = (d_flux1 + d_flux0)/2
            c = plt.contour(np.reshape(x_plasma, (129,129)), 
                            np.reshape(y_plasma, (129,129)), 
                            np.reshape(flux_plasma, (129,129)), 
                            levels=[d_flux])
            dp, cs = [], []
            for item in c.collections[0].get_paths():
                p = item.vertices
                dp.append((p[0,0]-p[-1,0])**2 + (p[0,1]-p[-1,1])**2)                
                cs.append(p)
                plt.plot(p[:,0], p[:,1], '.')
            
            d = min(dp)
            if d <  1e-5:
                d_flux0 = d_flux
            else:
                d_flux1 = d_flux
            
            index_d = np.argmin(dp)
            c_lcfs = cs[index_d]
            
        
        
        # find plasma inside lcfs
        rmin = min(c_lcfs[:, 0])
        rmax = max(c_lcfs[:, 0])
        zmin = min(c_lcfs[:, 1])
        zmax = max(c_lcfs[:, 1])
        self.fgs['xp'] = (rmin+rmax)/2
        
        index = flux_plasma*0
        index[flux_plasma>d_flux] = 1         
        index[x_plasma>rmax] = 0       
        index[x_plasma<rmin] = 0       
        index[y_plasma>zmax] = 0      
        index[y_plasma<zmin] = 0 
        
        # set plasma geo parameter
        self.fgs['Xp'] = (rmin+rmax)/2
        self.fgs['Yp'] = (zmin+zmax)/2
                
        flux_plasma_norm = (flux_plasma-max_flux_plasma)/(d_flux-max_flux_plasma)
        
        return flux_plasma, flux_plasma_norm, index
        # return flux_plasma_norm, index


    def find_X_point(self, flux_plasma, isplot=0):
        """
        find X point from flux of plasma, 

        input:
            option name         data_type    size     unit    meaning
            1      flux_plasma  array        (nr*nz,1)        flux of plasma region induced by plasma and pfs
            0      isplot       int          1                is plot or not. 


        output:
            option  name        data_type   size    unit    meaning
            1       X_point     list        n       m       coordinate of x points 

        reference:
            [1]

        Developed by xiangguchn 2023/12/25
        """
        
        X_point=[]
        flux_plasma = np.reshape(flux_plasma, (129,129))
        
        # x point up
        # nx  = self.fgs['index_xpoint_up_x']
        # ny  = self.fgs['index_xpoint_up_y']
        # dnx = self.fgs['index_xpoint_width']
        # dny = self.fgs['index_xpoint_height']
        # flux_plasma_up = flux_plasma[int(nx-dnx/2):int(nx+dnx/2), 
        #                              int(ny-dny/2):int(ny+dny/2)]
        
        flux_plasma_up = flux_plasma[-30:, 10:60]
        
        indexs_col_min = np.argmin(flux_plasma_up, axis=0)
        indexs_row_max = np.argmax(flux_plasma_up, axis=1)     
        
        for index in range(len(indexs_row_max)):
            if indexs_col_min[indexs_row_max[index]] == index:
                X_point.append([indexs_row_max[index], index])
        
        if isplot:
            plt.figure()
            plt.contour(flux_plasma_up,15)
            for index in range(len(indexs_row_max)):
                plt.plot(indexs_row_max[index], index, 'b.')
            
            for index in range(len(indexs_col_min)):
                plt.plot(index, indexs_col_min[index], 'k.')
            
            for x in X_point:
                plt.plot(x[0], x[1], 'r*')
        
        
        # x point down   
        flux_plasma_down = flux_plasma[0:30, 10:60]
        
        indexs_col_min = np.argmin(flux_plasma_down, axis=0)
        indexs_row_max = np.argmax(flux_plasma_down, axis=1)     
        
        for index in range(len(indexs_row_max)):
            if indexs_col_min[indexs_row_max[index]] == index:
                X_point.append([indexs_row_max[index], index])
        
        if isplot:
            plt.figure()
            plt.contour(flux_plasma_down,15)
            for index in range(len(indexs_row_max)):
                plt.plot(indexs_row_max[index], index, 'b.')
            
            for index in range(len(indexs_col_min)):
                plt.plot(index, indexs_col_min[index], 'k.')
            
            for x in X_point:
                plt.plot(x[0], x[1], 'r*')
        
        return X_point



    def identity_plasma(self, flux_plasma, flux_limiter, init_factor=1):
        """
        find LCFS of plasma region with flux_plasma and flux_limiter. 
        judge plasma is limiter or divertor configuration with inductance on limiter

        input:
            option name         data_type    size     unit    meaning
            1      flux_plasma  array        (nr*nz,1)        flux of plasma induced by plasma and pfs
            1      flux_limiter array        (nlimiter,1)     flux of limiter induced by plasma and pfs
            0      isplot       int          1                is plot or not. 

        output:
            option  name            data_type   size    unit    meaning
            1      flux_plasma      array       (nr*nz,1)       flux of plasma induced by plasma and pfs
            1      flux_plasma_norm array       (nr*nz,1)       flux of plasma normalized
            1      index            array       (nr*nz,1)       index of plasma: 1 means inside LCFS; 0 means outside LCFS
            
        reference:
            [1]

        Developed by xiangguchn 2023/12/25
        """
        
        
        # plasma and limiter region
        x_plasma = self.eq['green']['grid']['plasma']['x']
        y_plasma = self.eq['green']['grid']['plasma']['y']
        
        x_limiter = self.eq['machine']['limiter']['x_limiter']
        y_limiter = self.eq['machine']['limiter']['y_limiter']
        
        
        X_point = self.find_X_point(flux_plasma)
        
        
        
        v_step = 1e-3 * init_factor
        ok_step = 5e-6 * init_factor
        
        my_delay = 0.5
        
        y_top_limite = 1
        
        max_flux_plasma = max(flux_plasma[abs(y_plasma)<y_top_limite])
        index_max_flux_plasma = np.argmax(flux_plasma[abs(y_plasma)<y_top_limite])
        # plt.plot(x_plasma[index_max_flux_plasma], y_plasma[index_max_flux_plasma], 'r*')
        
        # position higher than 1m should be cancelled
        x_limiter_2 = x_limiter[abs(y_limiter)<y_top_limite]
        y_limiter_2 = y_limiter[abs(y_limiter)<y_top_limite]
        flux_limiter_2 = flux_limiter[abs(y_limiter)<y_top_limite]
        
        
        # find the max flux at limiter
        max_flux_limiter = max(flux_limiter_2)
        index_flux_limiter = np.argmax(flux_limiter_2)
        
        
        # find LCFS
        d_flux1 = max_flux_limiter
        d_flux0 = max_flux_plasma
        d_flux = d_flux1
        # plt.figure()
        # self.plot_flux(flux_plasma)
        while abs(d_flux0-d_flux1) > ok_step:
            d_flux = (d_flux1 + d_flux0)/2
            c = plt.contour(np.reshape(x_plasma, (129,129)), 
                            np.reshape(y_plasma, (129,129)), 
                            np.reshape(flux_plasma, (129,129)), 
                            levels=[d_flux])
            dp, cs = [], []
            for item in c.collections[0].get_paths():
                p = item.vertices
                dp.append((p[0,0]-p[-1,0])**2 + (p[0,1]-p[-1,1])**2)                
                cs.append(p)
                # plt.plot(p[:,0], p[:,1], '.')
            
            d = min(dp)
            if d <  1e-5:
                d_flux0 = d_flux
            else:
                d_flux1 = d_flux
            
            index_d = np.argmin(dp)
            c_lcfs = cs[index_d]
            print(d_flux)
            
        self.fgs['flux_LCFS'] = d_flux
        
        
        # # find LCFS
        # d_flux1 = max_flux_limiter
        # d_flux0 = max_flux_plasma
        # d_flux = d_flux1
        # fig = plt.figure()
        # plt.ion()
        
        # c = plt.contour(np.reshape(x_plasma, (129,129)), 
        #                 np.reshape(y_plasma, (129,129)), 
        #                 np.reshape(flux_plasma, (129,129)))
        
        # c = plt.contour(np.reshape(x_plasma, (129,129)), 
        #                 np.reshape(y_plasma, (129,129)), 
        #                 np.reshape(flux_plasma, (129,129)), 
        #                 levels=[d_flux1])
        
        # c = plt.contour(np.reshape(x_plasma, (129,129)), 
        #                 np.reshape(y_plasma, (129,129)), 
        #                 np.reshape(flux_plasma, (129,129)), 
        #                 levels=[d_flux0])
        
        # while abs(d_flux0-d_flux1) > ok_step:
            
        #     plt.title('abs(d_flux0-d_flux1)=' + str(round(abs(d_flux0-d_flux1),5)) + '>' + str(ok_step))
            
            
        #     d_flux = (d_flux1 + d_flux0)/2
        #     c = plt.contour(np.reshape(x_plasma, (129,129)), 
        #                     np.reshape(y_plasma, (129,129)), 
        #                     np.reshape(flux_plasma, (129,129)), 
        #                     levels=[d_flux])
        #     dp, cs = [], []
        #     for item in c.collections[0].get_paths():
        #         p = item.vertices
        #         dp.append((p[0,0]-p[-1,0])**2 + (p[0,1]-p[-1,1])**2)                
        #         cs.append(p)
        #         plt.plot(p[:,0], p[:,1], '.')
            
        #     d = min(dp)
        #     if d <  1e-5:
        #         d_flux0 = d_flux
        #     else:
        #         d_flux1 = d_flux
            
        #     index_d = np.argmin(dp)
        #     c_lcfs = cs[index_d]
        #     print(d_flux)
        #     plt.pause(0.5)
        
        # plt.title('abs(d_flux0-d_flux1)=' + str(round(abs(d_flux0-d_flux1),5)) + '<' + str(ok_step))
        
        # plt.ioff()
        # plt.show()
        
        
        
        # find plasma inside lcfs
        rmin = min(c_lcfs[:, 0])
        rmax = max(c_lcfs[:, 0])
        zmin = min(c_lcfs[:, 1])
        zmax = max(c_lcfs[:, 1])
        self.fgs['xp'] = (rmin+rmax)/2
        
        
        index = flux_plasma*0
        index[flux_plasma>d_flux] = 1         
        index[x_plasma>rmax] = 0       
        index[x_plasma<rmin] = 0       
        index[y_plasma>zmax] = 0      
        index[y_plasma<zmin] = 0 
        
        # set plasma geo parameter
        self.fgs['Xp'] = (rmin+rmax)/2
        self.fgs['Yp'] = (zmin+zmax)/2
        
        
        flux_plasma_norm = (flux_plasma-max_flux_plasma)/(d_flux-max_flux_plasma)
        # return flux_plasma_norm, index

        return flux_plasma, flux_plasma_norm, index

       
    
    
    def get_pprime_curent(self, flux_plasma_norm):
        """
        get plasma current pf pprime part with norm flux of plasma, 

        input:
            option name              data_type    size     unit    meaning
            1      flux_plasma_norm  array        (nr*nz,1)        flux of plasma region induced by plasma and pfs

        output:
            option  name            data_type   size      unit    meaning
            1       j_pprime        array       (nr*nz,1)         plasma distribution of pprime part

        reference:
            [1]

        Developed by xiangguchn 2023/12/25
        """
        
        m_p, n_p = 1, 1
        
        i = np.argmin(flux_plasma_norm)
        x_plasma = self.eq['green']['grid']['plasma']['x']
        betap0 = self.eq['machine']['plasma']['betap'] 
        
        # xp = self.fgs['Xp'] 
        xp = self.eq['green']['grid']['plasma']['x'][i]
        xp = self.fgs['xp']
        
        # flux_1 = flux_plasma
        # flux_2 = 1-flux_1
        
        j_pprime = betap0 * x_plasma/xp * (1-flux_plasma_norm**n_p)**m_p
        
        return j_pprime
    
    
    def get_ffrime_curent(self, flux_plasma_norm):
        """
        get plasma current pf ffrime part with norm flux of plasma, 

        input:
            option name              data_type    size     unit    meaning
            1      flux_plasma_norm  array        (nr*nz,1)        flux of plasma region induced by plasma and pfs

        output:
            option  name            data_type   size      unit    meaning
            1       j_ffrime        array       (nr*nz,1)         plasma distribution of ffrime part

        reference:
            [1]

        Developed by xiangguchn 2023/12/25
        """
        
        m_f, n_f = 1, 1
        
        x_plasma = self.eq['green']['grid']['plasma']['x']
        xp = self.fgs['xp']
        
        betap0 = self.eq['machine']['plasma']['betap'] 
        
        # flux_1 = flux_plasma/max(flux_plasma)
        # flux_2 = 1-flux_1
        
        j_ffrime = ((1-betap0) * xp/x_plasma) * (1-flux_plasma_norm**n_f)**m_f
        
        return j_ffrime
    
    
    
    def get_init_plasma_para(self, isplot=1):
        
        
        Ip = self.eq['machine']['plasma']['Ip']     
        xp = self.eq['machine']['plasma']['Xp']        
        yp = self.eq['machine']['plasma']['Yp']     
        
        ap = self.eq['machine']['plasma']['ap']
        chi = self.eq['machine']['plasma']['chi']             
        tri = self.eq['machine']['plasma']['tri']
        Rstep = self.eq['machine']['plasma']['Rstep']
             
        self.eq['green']['delta'] = 0.02
        self.eq['green']['Lsol'] = 0.03
        self.eq['green']['alpha'] = 0.7854
        delta = self.eq['green']['delta'] + 0.0 
        Lsol  = self.eq['green']['Lsol']             
        alpha = self.eq['green']['alpha']     
        
        ap = 0.2119
        xp = 1.1381
        chi = 1.0899
        
        factor = np.array([1])
        n_level = math.ceil(ap/Rstep)
        Rstep = ap/n_level
        
        xp0 = xp+delta # delta is caused by shafranov shift
        X = np.array([xp0])
        Y = np.array([yp])
        
        
        for i in range(1, n_level+1):
            di = alpha/2/i
            xpi = xp0 - i*delta/n_level
            
            theta = np.arange(di, 2*np.pi+di, di)
            theta2 = 2 * theta
            theta2 = np.where((theta2<=np.pi) | (theta2>=3*np.pi), theta2 , 0)
            
            r_plasma = i*(1-i/n_level*tri*np.abs(np.sin(theta2))**1.6/3) * Rstep
            
            Y = np.concatenate((Y, yp + chi*r_plasma*np.sin(theta)))
            
            com = np.array([np.abs(complex(np.cos(ta))**1.6) for ta in theta])
            X = np.concatenate((X, xpi + r_plasma*np.cos(theta) - 
                               i/n_level*r_plasma*(1-com)*tri))
        
            factor = np.concatenate((factor, np.ones(shape=(len(theta))) * 
                                     (1-(i*Rstep/(ap+Lsol))**2)**1.4))
            
        I = factor*Ip/np.sum(factor)     
        
        if isplot:
            plt.figure()
            plt.plot(X, Y, '.')
        
        
        
        return I, X, Y
        
    
    
    def calculate_mi_plasma_init(self, x_plasma, y_plasma, x2, y2, I_init, gapx=0):
        
        if gapx==0:
            gapx = (x_plasma[1]-x_plasma[0])/2
        
        R1 = np.sqrt((x_plasma+x2)**2 + (y_plasma-y2)**2);
        m = 4 * x2 * x_plasma / R1**2;

        # avoiding the overlay of source and field
        index = np.where(abs(m-1)<1e-10)
    
        
        xx1 = x_plasma; # may modify, not exist outside [if ~isempty(index)]
        if index is not None:
            xx1[index] = x_plasma[index] + gapx;    # move 1 mm outside
            R1 = np.sqrt((xx1+x2)**2 + (y_plasma-y2)**2)
            m  = 4 * x2 * xx1 / R1**2;
        
        
        myk = special.ellipk(m)
        mye = special.ellipe(m)
        mi_plasma_plasma = self.cmu * I_init * R1 * (2 * (myk-mye)-m * myk);
        
        return mi_plasma_plasma
    
    
    def calculate_mi_limiter_init(self, x_limiter, y_limiter, x_init, y_init, I_init, gapx):


        R1 = np.sqrt(((x_limiter+x_init)**2 + (y_limiter-y_init)**2).astype('float'));
        m = 4 * x_init * x_limiter / R1**2;

        # avoiding the overlay of source and field
        index = np.where(abs(m-1)<1e-10)
    
        
        xx1 = x_limiter; # may modify, not exist outside [if ~isempty(index)]
        if index is not None:
            xx1[index] = x_limiter[index] + gapx;    # move 1 mm outside
            R1 = np.sqrt(((xx1+x_init)**2 + (y_limiter-y_init)**2).astype('float'))
            m  = 4 * x_init * xx1 / R1**2;
        
        
        myk = special.ellipk(m.astype('float'))
        mye = special.ellipe(m.astype('float'))
        mi_limiter_init = self.cmu * I_init * R1 * (2 * (myk-mye)-m * myk);
        
        return mi_limiter_init
    
    
    
    def init_plasma(self, pfs, isplot=0):
        """
        initial plasma for intergration 
        
        """
        
        # green funciton 
        mi_plasma_pf = self.eq['green']['mi_plasma_pf']
        mi_limiter_pf = self.eq['green']['mi_limiter_pf']
        
        # plasma and limiter region
        x_plasma = self.eq['green']['grid']['plasma']['x']
        y_plasma = self.eq['green']['grid']['plasma']['y']
        num_plasma = len(x_plasma)
        gapx = (x_plasma[1]-x_plasma[0])/2
        
        
        
        # flux of plasma
        # flux inducted by plasma on plasma
        I_init, x_init, y_init = self.get_init_plasma_para()
        num_init = len(x_init)
        # flux_plasma_plasma = self.calculate_mi_plasma_init(x_plasma, y_plasma, 
        #                                                    x_init, y_init, I_init)
        num_init = len(I_init)
        
        flux_plasma_init = np.zeros(shape=(num_plasma, num_init))
        # pool = mp.Pool(self.num_cores)        
        # flux_plasma_init = [pool.apply(self.calculate_mi_plasma_init, 
        #                                args=(x_plasma, y_plasma, x_init[i], y_init[i], I_init[i], gapx)) 
        #                     for i in range(num_init)]        
        for i in range(num_init):
            flux_plasma_init[:,i] = self.calculate_mi_plasma_init(x_plasma, y_plasma, 
                                                                  x_init[i], y_init[i], 
                                                                  I_init[i], gapx)
        flux_plasma_init = np.sum(flux_plasma_init, axis=1)
        
        # flux inducted by PF current on plasma
        flux_plasma_pfs = pfs.dot(mi_plasma_pf)
        flux_plasma_pfs = flux_plasma_pfs.reshape(x_plasma.shape)
        
        # flux on plasma
        flux_plasma = flux_plasma_init + flux_plasma_pfs
        
        
        
        # flux of limiter
        x_limiter = self.eq['machine']['limiter']['x_limiter']
        y_limiter = self.eq['machine']['limiter']['y_limiter']
        num_limiter = len(x_limiter)
        
        # flux inducted on limiter by pf
        flux_limiter_pfs = pfs.dot(mi_limiter_pf)
        
        # flux inducted on limiter by init plasma        
        flux_limiter_init = np.zeros(shape=(num_limiter, num_init));
        for i in range(num_init):
            flux_limiter_init[:,i] = self.calculate_mi_limiter_init(x_limiter, y_limiter, 
                                                                  x_init[i], y_init[i], 
                                                                  I_init[i], gapx)
        flux_limiter_init = np.sum(flux_limiter_init, axis=1)
        
        # flux inducted on limiter
        flux_limiter = flux_limiter_pfs + flux_limiter_init
        
        
        
        # 3: identify plasma of init
        # identify plasma with plasma and limiter flux
        flux_plasma, flux_plasma_norm, index = self.identity_plasma_init(flux_plasma, flux_limiter, init_factor=10)
        
        
        
        # calculate plasma with phi
        j_plasma = self.iterate_plasma_with_phi(flux_plasma, flux_plasma_norm, index)        
        
        
        #
        if isplot:
            plt.figure()
            plt.plot(x_init, y_init, '.')
        
        return j_plasma, flux_plasma
    
    
    
    def iterate_plasma_with_phi(self, flux_plasma, flux_plasma_norm, index_plasma, cratio=1):
        
        
        
        x_plasma = self.eq['green']['grid']['plasma']['x']
        y_plasma = self.eq['green']['grid']['plasma']['y']
        
        # plt.figure()
        # plt.contour(np.reshape(flux_plasma,(129,129)))
        
        j_pprime = self.get_pprime_curent(flux_plasma_norm)
        j_pprime = index_plasma*j_pprime    # only region inside lcfs exsit plasma
        
        j_ffrime = self.get_ffrime_curent(flux_plasma_norm)
        j_ffrime = index_plasma*j_ffrime    # only region inside lcfs exsit plasma
        
        j_plasma = j_pprime + j_ffrime
        
        
        # sum of plasma current inside LCFS should equil to total plasma  
        j_plasma = j_plasma/sum(j_plasma) * plasma
        
        return j_plasma
    
    
    
    
    def iterate_phi_with_plasma(self, j_plasma):
                
        x_plasma = np.reshape(self.eq['green']['grid']['plasma']['x'], (129*129))
        y_plasma = np.reshape(self.eq['green']['grid']['plasma']['y'], (129*129))
        
        
        
        if 1:
            j_plasma=np.reshape(j_plasma, (129,129))
            j_plasma[:64,:] = (j_plasma[:64,:] + j_plasma[-1:-65:-1,:])/2
            j_plasma[-1:-65:-1,:] = j_plasma[:64,:]
        
            j_plasma=np.reshape(j_plasma, (129*129))
        
        
        
        
        # flux on plasma by pf and plasma
        flux_plasma_pfs = pfs.dot(self.eq['green']['mi_plasma_pf'])               # by pf 
        flux_plasma_pfs = flux_plasma_pfs.reshape(x_plasma.shape)
        flux_plasma_plasma = j_plasma.dot(self.eq['green']['mi_plasma_plasma'])   # by plasma
        flux_plasma_plasma = flux_plasma_plasma.reshape(x_plasma.shape)        
        flux_plasma = flux_plasma_plasma + flux_plasma_pfs
        
        
        # flux on limiter by pf and plasma
        flux_limiter_pfs = pfs.dot(self.eq['green']['mi_limiter_pf'])               # by pf 
        flux_limiter_plasma = j_plasma.dot(self.eq['green']['mi_limiter_plasma'])   # by plasma
        flux_limiter = flux_limiter_pfs + flux_limiter_plasma
        
                
        # identify plasma with plasma and limiter flux
        flux_plasma, flux_plasma_norm, index = self.identity_plasma(flux_plasma, flux_limiter, 10)
        
        
        
        return flux_plasma, flux_plasma_norm, index
    
    
    
    def free_boundary_iterate(self, plasma, pfs, n_iterate):
        
        j_plasma, flux_plasma = self.init_plasma(pfs)       
        
        for i in range(n_iterate):
            print(i)
            flux_plasma, flux_plasma_norm, index = self.iterate_phi_with_plasma(j_plasma)
            
            j_plasma = self.iterate_plasma_with_phi(flux_plasma, flux_plasma_norm, index)
        
            self.plot_flux(flux_plasma, islcfs=1)
            
        return flux_plasma, j_plasma
        
    
    
    def plot_flux(self, flux_plasma, islcfs=0):
        
        plt.figure()
        x_plasma = np.reshape(self.eq['green']['grid']['plasma']['x'],(129,129))
        y_plasma = np.reshape(self.eq['green']['grid']['plasma']['y'],(129,129))
        c = plt.contour(x_plasma, y_plasma, flux_plasma.reshape((129,129)), 33)
        plt.clabel(c, inline=True)
        plt.plot(self.eq['machine']['limiter']['x_limiter'], self.eq['machine']['limiter']['y_limiter'], 'k*')
        
        if islcfs:
            c = plt.contour(np.reshape(x_plasma, (129,129)), 
                            np.reshape(y_plasma, (129,129)), 
                            np.reshape(flux_plasma, (129,129)), 
                            levels=[self.fgs['flux_LCFS']])
        
        ax = plt.gca()
        ax.set_aspect(1)
        
        
if __name__ == '__main__':

    
    FB = free_boundary_gs()
    
    
    plasma = 500000
    # pfs = np.array([0, 3000, 3000, 10000, 10000, -1500, -1500, -2500, -2500, -7800, -7800, 0, 0, 0, 0])
    # pfs = np.array([0, 2500, 2500,  4500,  4500,  1500,  1500, -4200, -4200, -3200, -3200, 0, 0, 0, 0])
    # pfs = np.array([0, 11000, 11000,  12000,  12000,  6000,  6000, -10000, -10000, -12500, -12500, 0, 0, 0, 0])
    pfs = np.array([0, 3000, 3000, 3000, 3000, -1500, -1500, -2500, -2500, -6800, -6800, 0, 0, 0, 0])
    n_interate = 1
    
    flux_plasma, j_plasma = FB.free_boundary_iterate(plasma, pfs, n_interate)
    
    
    FB.plot_flux(flux_plasma, islcfs=1)